#pragma once

#include "net_common.hpp"
#include "net_tsdeque.hpp"

namespace dw
{
    namespace net
    {
        // Preclaim for the function call
        // server_interface::OnValidate()
        template <typename T>
        class server_interface;

        // a user-defined encrypt function
        uint32_t Scramble(uint32_t);

        template <typename T>
        class connection : 
            public std::enable_shared_from_this<connection<T>>
        {
        public:
            enum class owner
            {
                server,
                client
            };

        protected:
            // Each connection has a unique socket to a remote
            asio::ip::tcp::socket m_socket;

            // On server, we have serveral connections,
            // but we want them in the same context
            // This conetext is shared with the whole asio instance,
            // which provided by the client or the server interface
            asio::io_context &m_asioContext;

            tsdeque<message<T>> m_qMsgOut;

            // This queue holds all messages that have been received from
            // the remote side of this connection. Note it is a reference
            // as the owner of this connection is expected to provide this
            tsdeque<owned_message<T>> &m_qMsgIn;
            message<T> m_msgTempIn;

            // variables used for validation
            uint32_t m_nRandomToken;
            uint32_t m_nScrambledTokenClient;
            uint32_t m_nScrambledTokenServer;

            uint32_t id; // 
            owner m_nOwnerType; // server or client

        public:
            connection(
                owner parent,
                asio::io_context &asioContext,
                asio::ip::tcp::socket socket,
                tsdeque<owned_message<T>> &qIn
            ) : 
                m_asioContext(asioContext),
                m_socket(std::move(socket)),
                m_qMsgIn(qIn)
            {
                id = 0;
                m_nOwnerType = parent;

                m_nRandomToken = 0;
                m_nScrambledTokenClient = 0;
                m_nScrambledTokenServer = 0;
            }

            virtual ~connection()
            {
                // done
            }
        
        public:
            void ConnectToClient(dw::net::server_interface<T> *server, uint32_t uid = 0)
            {
                if (owner::server == m_nOwnerType)
				{
					if (m_socket.is_open())
					{
                        id = uid;
                        // Was: ReadHeader();

                        // When a client attempts to connect to server
						// it should first validate itself by solve a puzzle
						// offered by the server
						m_nRandomToken = uint32_t(std::chrono::system_clock::now().time_since_epoch().count());
						m_nScrambledTokenServer = Scramble(m_nRandomToken);

                        std::cout << "Token: "
                                  << std::hex << m_nRandomToken << " "
                                  << std::dec << m_nRandomToken << " "
                                  << std::hex << htonl(m_nRandomToken) << " "
                                  << std::dec << htonl(m_nRandomToken) << std::endl;

                        std::cout << "Server Result: "
                                  << std::hex << m_nScrambledTokenServer << " "
                                  << std::dec << m_nScrambledTokenServer << " "
                                  << std::hex << htonl(m_nScrambledTokenServer) << " "
                                  << std::dec << htonl(m_nScrambledTokenServer) << std::endl;

                        // Transfer to Network Endian
						m_nRandomToken = htonl(m_nRandomToken);
						asio::async_write(m_socket, asio::buffer(&m_nRandomToken, sizeof(uint32_t)),
							[this](std::error_code ec, std::size_t length)
							{
								if (ec)
								{
                                    std::cout << "[" << id << "]: Write Validation Failed" << std::endl;
                                    m_socket.close();
                                }
							}
						);

						// issue a task to wait for client's response for validation
						asio::async_read(m_socket, asio::buffer(&m_nScrambledTokenClient, sizeof(uint32_t)),
							[this, server](std::error_code ec, std::size_t length)
							{
								if (!ec)
								{
									// Transfer to Host Endian
									m_nScrambledTokenClient = ntohl(m_nScrambledTokenClient);

                                    std::cout << "Client Return: "
                                              << std::hex << m_nScrambledTokenClient << " " 
                                              << std::dec << m_nScrambledTokenClient << " "
                                              << std::hex << htonl(m_nScrambledTokenClient) << " " 
                                              << std::dec << htonl(m_nScrambledTokenClient) << std::endl;
                                    
                                    // Compare sent data to actual solution
									if (m_nScrambledTokenClient == m_nScrambledTokenServer)
									{
										// Client has provided valid solution, so allow it to connect properly
										std::cout << "[SERVER] Client <" << id << "> Validated" << std::endl;
										server->OnClientValidate(this->shared_from_this());

										// Sit waiting to receive data now
										ReadHeader();
									}
									else
									{
										// Client gave incorrect data, so disconnect
										std::cout << "[SERVER] Client <" << id << "> Rejected (Fail Validation)" << std::endl;
                                        m_socket.close();
									}
								}
								else
								{
									// Some bigger failure occured
									std::cout << "[SERVER] Client Disconnected (Read Validation)" << std::endl;
									m_socket.close();
								}
							}
						);
                    }
					else
					{
						std::cout << "[SERVER] Failed to accept client <"<< uid <<">" << std::endl;
						m_socket.close();
					}
				}
            }

            void ConnectToServer(const asio::ip::tcp::resolver::results_type &endpoints)
            {
                if (owner::client == m_nOwnerType)
				{
                    asio::async_connect(
                        m_socket, endpoints, 
						[this](std::error_code ec, asio::ip::tcp::endpoint endpoint)
						{
							if (!ec)
							{
								//Was: ReadHeader();

                                // First thing server will do is to validate that
								// the connection is from a valid client
								asio::async_read(m_socket, asio::buffer(&m_nRandomToken, sizeof(uint32_t)),
									[this](std::error_code ec, std::size_t length)
									{
										if (!ec)
										{
											// Transfer to Host Endian
											m_nRandomToken = ntohl(m_nRandomToken);

											// Connection is a client, so solve puzzle
											m_nScrambledTokenClient = Scramble(m_nRandomToken);

											// Transfer to Network Endian
											m_nScrambledTokenClient = htonl(m_nScrambledTokenClient)
												;
											// Write the result to server
											asio::async_write(m_socket, asio::buffer(&m_nScrambledTokenClient, sizeof(uint32_t)),
												[this](std::error_code ec, std::size_t length)
												{
													if (!ec)
													{
														// Validation data sent, clients should sit and wait
														// for a response (or a closure)
														ReadHeader();
													}
													else
													{
														std::cout << "Disconnected from Server (Fail Validation)" << std::endl;
														m_socket.close();
													}
												}
											);
										}
										else
										{
											// Some bigger failure occured
											std::cout << "Disconnected from Server (Read Validation)" << std::endl;
											m_socket.close();
										}
									}
								);
							}
							else
							{
								std::cout << "Failed to connect to server." << std::endl;
                                m_socket.close();
                            }
						}
					);
				}
            }

            void Disconnect()
            {
                if (IsConnected()) {
		    		asio::post(m_asioContext, [this]() { m_socket.close(); });
		    	}
            }

            bool IsConnected() const
            {
                return m_socket.is_open();
            }

            void Send(const message<T> &msg)
            {
                asio::post(
                    m_asioContext, 
					[this, msg]()
					{
						bool bWritingMsg = !m_qMsgOut.empty();
						m_qMsgOut.push_back(msg);
						if (!bWritingMsg) {
							WriteHeader();
						}
					}
				);
            }

        public:
            uint32_t GetID() const {
                return id;
            }

        private:
            //! ASYNC - Write Message Header
            void WriteHeader()
            {
                message<T> msgOut = m_qMsgOut.front();

                asio::async_write(
                    m_socket,
                    asio::buffer(&msgOut.header.NetOrder(), sizeof(message_header<T>)),
                    [this](std::error_code ec, std::size_t length)
                    {
                        if (!ec)
                        {
                            std::cout << m_qMsgOut.front().header << std::endl;
                            if (m_qMsgOut.front().body.size() > 0)
                            {
                                WriteBody();
                            }
                            else
                            {
                                m_qMsgOut.pop_front();

                                if (!m_qMsgOut.empty())
                                {
                                    WriteHeader();
                                }
                            }
                        }
                        else
                        {
                            std::cout << "[" << id << "] Write Header Fail." << std::endl;
                            m_socket.close();
                        }
                    });
            }

            //! ASYNC - Write Message body
            void WriteBody()
            {
                asio::async_write(
                    m_socket,
                    asio::buffer(m_qMsgOut.front().body.data(), m_qMsgOut.front().body.size()),
                    [this](std::error_code ec, std::size_t length)
                    {
                        if (!ec)
                        {
                            // std::cout << m_qMsgOut.front().body.data() << std::endl;
                            m_qMsgOut.pop_front();

                            if (!m_qMsgOut.empty())
                            {
                                WriteHeader();
                            }
                        }
                        else
                        {
                            std::cout << "[" << id << "] Write Body Fail." << std::endl;
                            m_socket.close();
                        }
                    }
                );
            }

            //! ASYNC - Read Message Header
			void ReadHeader()
			{
                asio::async_read(
                    m_socket,
                    asio::buffer(&m_msgTempIn.header, sizeof(message_header<T>)),
                    [this](std::error_code ec, std::size_t length)
                    {
                        m_msgTempIn.header = m_msgTempIn.header.HostOrder();

                        if (!ec)
                        {
                            std::cout << m_msgTempIn.header << std::endl;
                            if (m_msgTempIn.header.size > 0)
							{
                                m_msgTempIn.body.resize(m_msgTempIn.header.size);
                                ReadBody();
							}
							else
                            {
								AddToIncomingMsgQueue();
							}
						}
						else
						{
							std::cout << "[" << id << "] Read Header Fail." << std::endl;
							m_socket.close();
						}
                    });
            }

            //! ASYNC - Read Message Body
			void ReadBody()
			{
				asio::async_read(
					m_socket,
					asio::buffer(m_msgTempIn.body.data(), m_msgTempIn.body.size()),
					[this](std::error_code ec, std::size_t length)
					{
						if (!ec)
						{
							AddToIncomingMsgQueue();
						}
						else
						{
							std::cout << "[" << id << "] Read Body Fail." << std::endl;
							m_socket.close();
						}
					}
				);
			}

			void AddToIncomingMsgQueue()
			{
				if (owner::server == m_nOwnerType) {
					m_qMsgIn.push_back({ this->shared_from_this(), m_msgTempIn });
				} else {
					m_qMsgIn.push_back({ nullptr, m_msgTempIn });
				}

				ReadHeader();
			}
        };
    }
}